# features.py
import re
import tldextract

# 初始化 extractor（禁用缓存）
tld_extractor = tldextract.TLDExtract(cache_dir=False)

def extract_domain_features(domain):
    # 使用 tldextract 提取域名信息
    ext = tld_extractor(domain)
    subdomain = ext.subdomain
    domain_name = ext.domain
    suffix = ext.suffix

    # 构造完整的域名
    full_domain = '.'.join([part for part in [subdomain, domain_name, suffix] if part])

    # 构造特征
    features = {
        "domain_length": len(full_domain),
        "num_dots": full_domain.count('.'),
        "num_digits": sum(c.isdigit() for c in full_domain),
        "num_hyphens": full_domain.count('-'),
        "num_vowels": sum(c in 'aeiou' for c in full_domain.lower()),
        "contains_ip": int(bool(re.match(r"^(\d{1,3}\.){3}\d{1,3}$", domain_name))),
        "starts_with_www": int(domain.startswith("www.")),
        "tld": suffix  # 将 TLD 作为特征返回
    }

    return features
